'''
predict.py有几个注意点
1、无法进行批量预测，如果想要批量预测，可以利用os.listdir()遍历文件夹，利用Image.open打开图片文件进行预测。
2、如果想要保存，利用r_image.save("img.jpg")即可保存。
3、如果想要原图和分割图不混合，可以把blend参数设置成False。
4、如果想根据mask获取对应的区域，可以参考detect_image中，利用预测结果绘图的部分。
seg_img = np.zeros((np.shape(pr)[0],np.shape(pr)[1],3))
for c in range(self.num_classes):
    seg_img[:, :, 0] += ((pr == c)*( self.colors[c][0] )).astype('uint8')
    seg_img[:, :, 1] += ((pr == c)*( self.colors[c][1] )).astype('uint8')
    seg_img[:, :, 2] += ((pr == c)*( self.colors[c][2] )).astype('uint8')
'''
import traceback
from Edge_Fitter import *
from progress.bar import IncrementalBar
from FPS_test import *

blend_path = os.path.join(results, 'blend-original')
compare_path = os.path.join(results, 'compare')
masks = os.path.join(python_file, 'mask')

def Img2cv2(pic):
    return cv2.cvtColor(np.array(pic), cv2.COLOR_RGB2BGR)

def main(pic, net):
    try:
        image = Image.open(pic)
        blend, mask = net.detect_image(image)
        if not blend is None:
            blend.save(pic.replace(original, blend_path))
            try_path = pic.replace(original, masks).replace('jpg', 'png')
            if os.path.exists(try_path):
                original_mask = Image.fromarray(cv2.imread(try_path))
                Image.blend(blend, original_mask, 0.3).save(pic.replace(original, compare_path))
        mask.save(pic.replace(original, masks_path))
        get_egg(Img2cv2(mask), np.array(image), pic.replace(original, results_path), pic)
    except:
        print('{:*^60}'.format('ERROR'))
        print('Open Error: ' + pic)
        exstr = traceback.format_exc()
        print(exstr)
        print('{:*^60}'.format(' '))

if __name__ == '__main__':
    start = time.time()
    pics = glob.glob(os.path.join(original, '*.jpg'))
    make_dirs([results, masks_path, results_path, blend_path, compare_path])
    bar = IncrementalBar(max=len(pics))
    Net = Pspnet()
    '''
    while True:
        img = input('Input image filename:')
        try:
            image = Image.open(img)
        except:         
            print('Open Error! Try again!')
            continue
        else:
            r_image = pspnet.detect_image(image)
            r_image.save(time.strftime('%Y_%m_%d_H_%M_%S.jpg'))
            r_image.show()
    '''
    for pic in pics:
        main(pic, Net)
        bar.next()
    end = time.time()
    print('\nTIME: ', end - start)
    FPS_test()